'''
如果设备有GPU，则调用GPU进行计算
在MacOS-M系列芯片上用MPS进行深度学习
'''
# 导入PyTorch深度学习框架的核心模块，包含张量操作、神经网络等基础功能
import torch
# 导入PyTorch的函数式接口模块，通常缩写为F，提供激活函数、池化等操作
import torch.nn.functional as F
# 从torchvision导入transforms模块，用于图像预处理（如格式转换、归一化等）
from torchvision import transforms
# 从torchvision导入datasets模块，包含常用的公开数据集（本例使用MNIST数据集等）
from torchvision import datasets
# 从torch.utils.data导入DataLoader，用于批量加载数据并支持多线程
from torch.utils.data import DataLoader
# 导入PyTorch的优化器模块，包含各种优化算法（如SGD、Adam等）
import torch.optim as optim
# 导入时间模块
import time
# 导入matplotlib来绘图
import matplotlib.pyplot as plt

# 检查可用的加速设备，优先使用MPS（Apple Silicon），其次是CUDA，最后用CPU
if torch.backends.mps.is_available():
    # 若Mac搭载Apple Silicon芯片且MPS可用（需macOS 12.3+及以上）
    device = torch.device("mps")
    print("使用MPS设备加速（Apple Silicon GPU）:", device)
elif torch.cuda.is_available():
    # 若有NVIDIA GPU且CUDA可用
    device = torch.device("cuda:0")
    print("使用CUDA设备加速:", device)
else:
    # 若无任何GPU加速，使用CPU
    device = torch.device("cpu")
    print("使用CPU运行:", device)


# 第一步：准备数据（数据加载与预处理）
# 定义批量大小（每次训练/测试输入模型的样本数量）
# 批量大小影响训练效率和模型收敛，64是常用的适中值
batch_size = 64

# 定义图像预处理流水线（Compose将多个预处理操作组合成一个流程）
transform = transforms.Compose([
    transforms.ToTensor(),  # 将PIL图像（0-255的像素值）转换为PyTorch张量（0-1的浮点数）
    # 对张量进行归一化：output = (input - mean) / std
    # MNIST数据集的均值为0.1307，标准差为0.3081（预先统计好的常用值）
    # 归一化可加速模型收敛，使不同特征在同一量级
    transforms.Normalize((0.1307, ), (0.3081, ))
])

# 加载MNIST训练数据集
# root：数据集保存路径（../dataset/mnist/表示上级目录的dataset/mnist文件夹）
# train=True：加载训练集（包含60000张手写数字图像）
# download=True：若本地没有数据集，则自动从网络下载
# transform=transform：应用上面定义的预处理流程
train_dataset = datasets.MNIST(root='../dataset/mnist/',train=True,download=True,transform=transform)

# 加载MNIST测试数据集
# train=False：加载测试集（包含10000张手写数字图像，用于评估模型性能）
test_dataset = datasets.MNIST(root='../dataset/mnist/',train=False,download=True,transform=transform)

# 创建训练数据加载器
# 将train_dataset按batch_size分成多个批次，shuffle=True表示每个epoch前打乱数据顺序
# 作用：打乱数据可避免模型学习到数据顺序的无关特征，提高泛化能力
train_loader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size)

# 创建测试数据加载器
# 测试时无需打乱数据（shuffle=False），按原顺序加载即可
test_loader = DataLoader(test_dataset, shuffle=False, batch_size=batch_size)



# 第二步：构建模型（定义CNN网络结构）
# 定义一个基于CNN（卷积神经网络）的模型类，继承自torch.nn.Module（PyTorch中所有模型的基类）
class cnnNet(torch.nn.Module):
    # 类的初始化方法，用于定义模型的层结构（只在实例化时执行一次）
    def __init__(self):
        # 调用父类（torch.nn.Module）的初始化方法，确保模型能正确注册参数和子模块
        super(cnnNet, self).__init__()

        # 定义第一个卷积层（Conv2d用于处理2D图像数据）
        # in_channels=1：输入图像的通道数（MNIST是灰度图，通道数为1；彩色图如RGB为3）
        # out_channels=10：卷积核的数量（输出特征图的通道数），10个卷积核提取10种基础特征
        # kernel_size=5：卷积核的尺寸为5x5（滑动窗口大小）
        self.conv1 = torch.nn.Conv2d(in_channels=1, out_channels=10, kernel_size=5)

        # 定义第二个卷积层（更深层次的特征提取）
        # in_channels=10：输入通道数需与上一层的输出通道数一致（10）
        # out_channels=20：输出通道数为20，提取更复杂的组合特征
        # kernel_size=5：仍使用5x5卷积核
        self.conv2 = torch.nn.Conv2d(in_channels=10, out_channels=20, kernel_size=5)

        # 定义最大池化层（MaxPool2d用于降维）
        # kernel_size=2：池化窗口大小为2x2，对窗口内的像素取最大值
        # 作用：减少特征图尺寸（缩小为原来的1/2），降低计算量，同时增强特征的平移不变性
        self.pooling = torch.nn.MaxPool2d(2)

        # 定义全连接层（Linear层用于分类决策）
        # in_features=320：输入特征的维度（由卷积层输出展平后得到，计算过程见forward方法）
        # out_features=10：输出维度为10（对应MNIST的10个数字类别）
        self.fc = torch.nn.Linear(320, 10)

    # 定义模型的前向传播逻辑（必须实现，PyTorch通过该方法计算输出）
    # x：输入的批量图像数据，形状为[batch_size, channels, height, width]
    # 例如MNIST输入为[64, 1, 28, 28]（64个样本，1通道，28x28像素）
    def forward(self, x):
        # 获取输入数据的批量大小（即一次输入的样本数量，如64）
        batch_size = x.size(0)

        # 第一层：卷积 -> 池化 -> 激活
        # 1. self.conv1(x)：对输入x应用第一个卷积层
        #    输入形状[64, 1, 28, 28] -> 输出形状[64, 10, 24, 24]
        #    计算方式：(28 - 5 + 1) = 24（卷积后尺寸=输入尺寸-卷积核尺寸+1）
        # 2. self.pooling(...)：对卷积结果做2x2最大池化
        #    输出形状[64, 10, 12, 12]（24/2=12，池化后尺寸=输入尺寸/池化窗口大小）
        # 3. F.relu(...)：应用ReLU激活函数，将负数置为0，引入非线性
        #    作用：使模型能拟合复杂的非线性关系，输出形状不变
        x = F.relu(self.pooling(self.conv1(x)))

        # 第二层：卷积 -> 池化 -> 激活
        # 1. self.conv2(x)：对上层输出应用第二个卷积层
        #    输入形状[64, 10, 12, 12] -> 输出形状[64, 20, 8, 8]
        #    计算方式：(12 - 5 + 1) = 8
        # 2. self.pooling(...)：2x2池化后，尺寸8/2=4
        #    输出形状[64, 20, 4, 4]
        # 3. F.relu(...)：ReLU激活，形状不变
        x = F.relu(self.pooling(self.conv2(x)))

        # 将卷积层输出的多维特征图展平为一维向量（全连接层要求输入为一维）
        # x.view(batch_size, -1)：将x重塑为[batch_size, 特征数]
        # -1表示自动计算特征数：20（通道）*4（高）*4（宽）=320
        # 输出形状[64, 320]
        x = x.view(batch_size, -1)

        # 全连接层计算：将320维特征映射到10维输出
        # 输出形状[64, 10]，每个维度对应一个类别的预测分数（未归一化的概率）
        x = self.fc(x)

        # 返回最终的预测结果
        return x


# 实例化定义好的CNN模型
model = cnnNet()

# 设备选择：优先使用MPS（Apple Silicon GPU），其次是CUDA，最后用CPU
# 适配Mac的Apple Silicon芯片（M1/M2/M3等）
device = torch.device(
    "mps" if torch.backends.mps.is_available()  # 检查MPS是否可用
    else "cuda:0" if torch.cuda.is_available()  # 检查CUDA是否可用
    else "cpu"  # 都不可用时使用CPU
)

# 将模型移动到目标设备（MPS/CUDA/CPU）
model.to(device)


# 第三步：构建损失函数和优化器
# 定义损失函数：交叉熵损失（CrossEntropyLoss）
# 适用于多分类任务，内部已包含SoftMax激活和负对数似然损失，直接输入未归一化的分数即可
criterion = torch.nn.CrossEntropyLoss()

# 定义优化器：随机梯度下降（SGD）
# model.parameters()：需要优化的模型参数（权重和偏置）
# lr=0.01：学习率（控制参数更新的步长，过大可能不收敛，过小收敛慢）
# momentum=0.5：动量参数（模拟物理中的动量，加速收敛并减少震荡）
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)


# 定义训练函数（每个epoch的训练过程）
def train(epoch):
    # 新增：将模型切换为训练模式（关键！确保训练层逻辑正常）
    model.train()
    # 初始化累计损失（用于打印训练进度）
    running_loss = 0.0
    epoch_start_time = time.time()  # 记录当前epoch训练开始时间
    # 遍历训练数据加载器中的每个批次
    # batch_idx：批次索引，data：每个批次的数据（包含输入和标签）
    for batch_idx, data in enumerate(train_loader, 0):
        # 从data中分离输入图像（inputs）和对应的标签（target）
        inputs, target = data
        # 将输入和标签移动到目标设备（与模型同设备）
        inputs, target = inputs.to(device), target.to(device)

        # 清空优化器的梯度（PyTorch会累积梯度，每次迭代前需清零）
        optimizer.zero_grad()

        # 前向传播：将输入传入模型，得到预测输出
        outputs = model(inputs)
        # 计算损失：预测输出与真实标签的差异
        loss = criterion(outputs, target)
        # 反向传播：计算损失对各参数的梯度（自动求导）
        loss.backward()
        # 优化器更新参数：根据梯度调整模型权重
        optimizer.step()

        # 累加当前批次的损失值（.item()将张量转换为Python数值）
        running_loss += loss.item()
        # 每300个批次打印一次平均损失（监控训练进度）
        if batch_idx % 300 == 299:
            # 打印信息：当前 epoch+1（从1开始）、批次索引+1、平均损失（除以2000是为了缩放数值）
            print('[%d, %5d] loss: %.3f' % (epoch + 1, batch_idx + 1, running_loss / 2000))
            # 重置累计损失，为下一轮累计做准备
            running_loss = 0.0
    # 新增：打印当前epoch总训练耗时
    epoch_end_time = time.time()
    epoch_total_time = epoch_end_time - epoch_start_time
    print(f'Epoch {epoch+1} 训练耗时: {epoch_total_time:.2f}秒')


# 定义测试函数（评估模型在测试集上的性能）
def test():
    # 新增：将模型切换为评估模式（关键！确保测试层逻辑正常）
    model.eval()
    # 初始化正确预测的样本数和总样本数
    correct = 0
    total = 0
    test_start_time = time.time()  # 记录测试开始时间
    # with torch.no_grad()：禁用梯度计算（测试时不需要求导，节省内存和计算资源）
    with torch.no_grad():
        # 遍历测试数据加载器中的每个批次
        for data in test_loader:
            # 分离输入图像和标签，并移动到目标设备
            inputs, target = data
            inputs, target = inputs.to(device), target.to(device)
            # 前向传播：得到预测输出
            outputs = model(inputs)
            # 从输出中获取预测结果：torch.max返回最大值和对应的索引
            # dim=1：沿第1维（类别维度）取最大值，predicted为预测的类别索引
            _, predicted = torch.max(outputs.data, dim=1)
            # 累加总样本数（当前批次的样本数）
            total += target.size(0)
            # 累加正确预测的样本数（预测类别与真实标签一致的数量）
            correct += (predicted == target).sum().item()
    accuracy = 100 * correct / total  # 计算准确率（百分比）
    # 计算并打印测试总耗时
    test_end_time = time.time()
    test_total_time = test_end_time - test_start_time
    # 打印测试集准确率：（正确数/总数）*100%，并显示具体的正确数和总数
    print('测试集精度: %d %% [%d/%d] ｜ 测试总耗时: %.2f秒\n' %
          (100 * correct / total, correct, total ,test_total_time))
    return accuracy  # 返回准确率，用于绘图


# 主程序入口：当脚本直接运行时执行以下代码
if __name__ == '__main__':
    print("模型开始执行")
    print(f"当前训练设备: {device}")
    overall_start_time = time.time()  # 记录整体训练开始时间

    # 创建列表存储每个epoch的测试准确率
    test_accuracies = []  # 存储test()返回的accuracy
    epochs = []           # 存储对应的epoch序号（1,2,...,10）

    # 训练10个epoch
    for epoch in range(10):
        # 执行训练函数，训练一个epoch
        train(epoch)
        # 每个epoch训练结束后，在测试集上评估一次模型性能
        acc=test()
        # test返回的acc结果添加到test_accuracies列表
        test_accuracies.append(acc)
        # epoch从1开始计数
        epochs.append(epoch + 1)

    # 新增：打印整体训练+测试总耗时
    overall_end_time = time.time()
    overall_total_time = overall_end_time - overall_start_time
    print(f'全部10个Epoch训练+测试总耗时: {overall_total_time:.2f}秒')

    # 训练结束后绘制准确率折线图
    # 设置支持中文的字体（解决中文显示警告）
    plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei', 'WenQuanYi Zen Hei']  # 适配Mac/Windows/Linux
    plt.rcParams['axes.unicode_minus'] = False  # 避免负号显示异常
    plt.figure(figsize=(10, 6))
    # 绘制折线图（x轴：epoch，y轴：准确率）
    plt.plot(epochs, test_accuracies, 'o-', color='darkorange', linewidth=2, markersize=8)
    # 添加标题和坐标轴标签
    plt.title('测试集准确率随训练轮次的变化', fontsize=15)
    plt.xlabel('训练轮次（Epoch）', fontsize=12)
    plt.ylabel('准确率（%）', fontsize=12)
    # 固定y轴范围（0-100%）
    plt.ylim(80, 100)
    # 添加网格线，方便观察数值
    plt.grid(linestyle='--', alpha=0.7)
    # 在每个点上标注具体准确率
    for i in range(len(epochs)):
        plt.text(epochs[i], test_accuracies[i] + 1, f'{test_accuracies[i]:.1f}%',
                 ha='center', fontsize=10)
    # 调整布局
    plt.tight_layout()
    # 显示图像
    plt.show()